import pandas as pd
from sklearn.model_selection import GroupKFold
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error
import shap
import matplotlib.pyplot as plt
import os
import joblib

# 读取数据
data_path = 'D:\\学习&科研\\华为手表项目\\华为数据\\试验记录表\\all_stages_df_statistics.csv'
df = pd.read_csv(data_path)

# 筛选出 state 为 'running' 的行
df = df[df['state'] == 'running']

# 将 polar_hr 和 polar_rr 列转换为适合模型的格式，例如取平均值
df['polar_hr'] = df['polar_hr'].apply(lambda x: eval(x)[0] if isinstance(eval(x), list) and len(eval(x)) > 0 else 0)
df['polar_rr'] = df['polar_rr'].apply(lambda x: eval(x)[0] if isinstance(eval(x), list) and len(eval(x)) > 0 else 0)

# 定义输入特征和目标变量
X = df[['speed', 'polar_hr_mean', 'polar_hr_min', 'polar_hr_max', 'polar_hr_median',
         'polar_hr_q1', 'polar_hr_q3', 'polar_rr_mean', 'polar_rr_median',
         'polar_rr_q1', 'polar_rr_q3', 'sex', 'age', 'hight', 'weight']]
y = df['physiology_RPE']
groups = df['number']  # 分组列

# 使用 GroupKFold 进行分组划分
group_kfold = GroupKFold(n_splits=10)

# 初始化结果列表
results = []
all_y_true = []
all_y_pred = []

# 创建保存SHAP图的文件夹
shap_plot_dir = 'shap_plots'
os.makedirs(shap_plot_dir, exist_ok=True)

# 基于分组进行划分和训练
for fold, (train_index, test_index) in enumerate(group_kfold.split(X, y, groups=groups)):
    X_train, X_test = X.iloc[train_index], X.iloc[test_index]
    y_train, y_test = y.iloc[train_index], y.iloc[test_index]
    group_test = groups.iloc[test_index]

    # 创建线性回归模型
    model = LinearRegression()

    # 训练模型
    model.fit(X_train, y_train)

    # 预测
    y_pred = model.predict(X_test)

    # 保存所有的真实值和预测值
    all_y_true.extend(y_test)
    all_y_pred.extend(y_pred)

    # 绘制每个 number 的实际值和预测值图
    for number in group_test.unique():
        plt.figure(figsize=(10, 6))
        mask = group_test == number  # 选择当前 number 的数据
        plt.scatter(y_test[mask], y_pred[mask], label=f'Number {number}', alpha=0.6)
        plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'k--', lw=2)  # y=x线
        plt.title(f'Actual vs Predicted for Number {number}')
        plt.xlabel('Actual Values')
        plt.ylabel('Predicted Values')
        plt.legend()
        plt.grid()
        plt.savefig(f'actual_vs_predicted_number_{number}_fold_{fold + 1}.png')  # 保存图像
        plt.close()

# 计算整体模型性能
overall_mse = mean_squared_error(all_y_true, all_y_pred)
overall_r2 = r2_score(all_y_true, all_y_pred)
overall_mae = mean_absolute_error(all_y_true, all_y_pred)

# 输出整体模型性能
print(f'Overall Mean Squared Error: {overall_mse}')
print(f'Overall R² Score: {overall_r2}')
print(f'Overall Mean Absolute Error: {overall_mae}')

# 使用最佳模型在整个训练集上进行最终训练
final_model = LinearRegression()
final_model.fit(X, y)

# 保存最终模型
joblib.dump(final_model, 'final_model.joblib')  # 保存模型

# 加载模型并进行预测
loaded_model = joblib.load('final_model.joblib')

# 进行预测（使用整个数据集或新的数据集）
predictions = loaded_model.predict(X)

# 创建一个 DataFrame 来保存真实值和预测值
comparison_df = pd.DataFrame({
    'True Values': y,
    'Predictions': predictions,
    'number': df['number']
})

# 按 number 保存每个 group's 预测结果
for number in comparison_df['number'].unique():
    group_df = comparison_df[comparison_df['number'] == number]
    group_df.to_csv(f'predictions_comparison_number_{number}.csv', index=False)

# 打印预测结果
print("预测结果和真实值比较：")
print(comparison_df)
